from transformers import AutoModelForCausalLM, AutoTokenizer

# 指定本地模型目录路径（包含所有文件的文件夹）
model_dir = "D:\model"

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    attn_implementation="eager",  # 禁用flash-attention
    device_map="auto",  # 自动分配CPU/GPU
    trust_remote_code=True,  # 允许执行自定义代码（针对phi-3特殊结构）
    torch_dtype="auto"       # 自动选择数据类型
)

tokenizer = AutoTokenizer.from_pretrained(
    model_dir,
    trust_remote_code=True
)

# 测试推理
prompt = "北京和上海的区别是什么？"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=100)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))